package com.one.blocks.eds.core;

import com.one.blocks.eds.config.properties.EdsProperties;
import com.one.blocks.eds.config.properties.EdsRemoteProperties;
import com.one.blocks.eds.exception.EdsException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.io.File;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author <a href="mailto:idler41@163.con">linfuxin</a> created on 2023-09-11 10:03:42
 */
@Slf4j
public class ExpireDataSource extends AbstractRoutingDataSource implements DisposableBean {

    final Map<String, ExpireDataSourceHolder> staticDataSources = new ConcurrentHashMap<>(8);

    final Map<String, ExpireDataSourceHolder> targetDataSources = new ConcurrentHashMap<>(32);

    final BlockingQueue<ExpireDataSourceHolder> evictQueue = new LinkedBlockingQueue<>();

    /**
     * key => jarUrl, value => UrlClassLoader
     */
    private final Map<String, ClassLoader> classLoaderMap = new ConcurrentHashMap<>(32);

    private final EdsPropertiesLoader edsPropertiesLoader;

    private final EdsProperties staticDataSourceProperties;

    private static final ThreadLocal<String> LAST_DS_KEY_WHEN_REGISTRY = new ThreadLocal<>();

    private final AtomicBoolean initialized = new AtomicBoolean(false);

    private DestroyEdsThread destroyThread;

    private ScanExpiredEdsThread scanThread;

    public ExpireDataSource(EdsPropertiesLoader edsPropertiesLoader, EdsProperties staticDataSourceProperties) {
        this.edsPropertiesLoader = edsPropertiesLoader;
        this.staticDataSourceProperties = staticDataSourceProperties;
    }

    @Override
    public void afterPropertiesSet() {
        if (!initialized.compareAndSet(false, true)) {
            // 已经初始化过，直接返回
            return;
        }

        if (!CollectionUtils.isEmpty(staticDataSourceProperties.getEds())) {
            // 注册永不过期数据源
            staticDataSourceProperties.getEds().forEach(this::registerStatic);
        }

        // 把线程启动放到初始化静态数据源后面，避免应用启动后马上失败关闭容器，2个线程没有被检测到关闭导致内存泄漏警告
        destroyThread = new DestroyEdsThread(staticDataSourceProperties.getEdsDestroyTaskName());
        scanThread = new ScanExpiredEdsThread(staticDataSourceProperties.getEdsScanTaskName());
        destroyThread.start();
        scanThread.start();

    }


    @Override
    public void destroy() throws Exception {
        // 中断并等待线程结束
        if (destroyThread != null) {
            destroyThread.interrupt();
        }
        if (scanThread != null) {
            scanThread.interrupt();
        }

        try {
            if (destroyThread != null) {
                destroyThread.join(5000);
            }
            if (scanThread != null) {
                scanThread.join(5000);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        // 关闭所有数据源
        closeAllDataSources();
    }

    @Override
    protected DataSource determineTargetDataSource() {
        String lookupKey = (String) determineCurrentLookupKey();
        lookupKey = lookupKey != null ? lookupKey : staticDataSourceProperties.getEdsDefaultKey();
        if (lookupKey == null) {
            throw new EdsException("both local and default datasource key are null!");
        }
        if (log.isDebugEnabled()) {
            log.debug("switch datasource key={}", lookupKey);
        }
        if (staticDataSources.containsKey(lookupKey)) {
            return staticDataSources.get(lookupKey).getDataSource();
        }

        ExpireDataSourceHolder dsHolder = targetDataSources.get(lookupKey);
        // 即使判断存在并发问题导致获取到旧数据源，也等同于加堵塞锁之前获取到数据源，处理逻辑一致
        if (dsHolder == null || dsHolder.isRemoved()) {
            dsHolder = register(lookupKey);
        }

        dsHolder.setAccessTime(System.currentTimeMillis());
        return dsHolder.getDataSource();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return LocalEdsRoutingKeyHolder.getKey();
    }

    /**
     * 指定驱动包创建数据源，通过edsPropertiesLoader加载驱动及其他配置信息
     *
     * @param dsKey 数据源唯一标识
     * @return 数据源引用
     */
    protected ExpireDataSourceHolder createDataSource(String dsKey) {
        if (log.isDebugEnabled()) {
            log.debug("register dataSource!key={}", dsKey);
        }
        EdsRemoteProperties edsRemoteProperties = edsPropertiesLoader.load(dsKey);
        if (edsRemoteProperties == null) {
            throw new IllegalStateException("eds初始化配置不能为null!dsKey=" + dsKey);
        }
        return createDataSource(dsKey, edsRemoteProperties);
    }

    protected ClassLoader getClassLoaderFromCache(EdsRemoteProperties edsRemoteProperties, ClassLoader parentClassLoader) {
        // 加载指定jar包，或指定路径下所有jar包
        URL[] urls = resolveUrls(edsRemoteProperties.getJarPath());
        return classLoaderMap.computeIfAbsent(edsRemoteProperties.getJarPath(), p -> new URLClassLoader(urls, parentClassLoader) {
            @Override
            public Class<?> loadClass(String name) throws ClassNotFoundException {
                try {
                    // 先检查是否已经加载了该类
                    Class<?> loadedClass = findLoadedClass(name);
                    // 再尝试从当前loader加载
                    return loadedClass != null ? loadedClass : findClass(name);
                } catch (ClassNotFoundException e) {
                    return super.loadClass(name);
                }
            }
        });
    }

    protected void registerStatic(String dsKey, Properties properties) {
        staticDataSources.computeIfAbsent(dsKey, k -> {
            if (log.isDebugEnabled()) {
                log.debug("register static dataSource!key={}", dsKey);
            }
            // 指定路径加载驱动
            if (properties.containsKey("edsDriverJarUrl")) {
                String jarUrl = properties.getProperty("edsDriverJarUrl");
                // 删除不属于HikariConfig的配置
                properties.remove("edsDriverJarUrl");
                EdsRemoteProperties edsRemoteProperties = EdsRemoteProperties.builder()
                        .jarPath(jarUrl)
                        .hikariProperties(properties)
                        .build();
                return createDataSource(dsKey, edsRemoteProperties);
            }
            return new ExpireDataSourceHolder(dsKey, properties);
        });
    }

    protected ExpireDataSourceHolder register(String dsKey) {
        try {
            // createDataSource内部的edsPropertiesLoader.load实现如果也要访问数据源，且与当前dsKey相同，则会存在死锁
            String lastDsKey = LAST_DS_KEY_WHEN_REGISTRY.get();
            if (lastDsKey != null && lastDsKey.equals(dsKey)) {
                throw new EdsException("发现死锁!请切换到正确数据源，当前数据源=" + dsKey);
            }
            LAST_DS_KEY_WHEN_REGISTRY.set(dsKey);
            return targetDataSources.computeIfAbsent(dsKey, k -> createDataSource(dsKey));
        } finally {
            LAST_DS_KEY_WHEN_REGISTRY.remove();
        }
    }

    /**
     * 指定驱动包创建数据源，edsRemoteProperties指定了驱动包位置
     *
     * @param dsKey      数据源key
     * @param properties 包含驱动包位置的数据源配置信息
     * @return 数据源
     * @see com.zaxxer.hikari.util.DriverDataSource Enumeration<Driver> drivers = DriverManager.getDrivers();
     */
    protected ExpireDataSourceHolder createDataSource(String dsKey, EdsRemoteProperties properties) {
        ClassLoader oldClassLoader = Thread.currentThread().getContextClassLoader();
        ClassLoader driverClassLoader = getClassLoaderFromCache(properties, oldClassLoader);
        try {
            // 提前设置 ClassLoader，确保 DriverManager 使用正确的 ClassLoader
            Thread.currentThread().setContextClassLoader(driverClassLoader);
            return buildExpireDataSourceHolder(dsKey, properties.getHikariProperties());
        } finally {
            Thread.currentThread().setContextClassLoader(oldClassLoader);
        }
    }

    protected ExpireDataSourceHolder buildExpireDataSourceHolder(String dsKey, Properties properties) {
        return new ExpireDataSourceHolder(dsKey, properties);
    }

    /**
     * 将jar文件或指定路径下所有jar文件转为URL数组
     * @param jarPath 既可以是jar文件，也可以是目录
     * @return jar文件
     */
    private URL[] resolveUrls(String jarPath) {
        File pathFile = new File(jarPath);
        if (pathFile.isDirectory()) {
            // 如果是目录，加载目录下所有jar文件
            File[] jarFiles = pathFile.listFiles((dir, name) -> name.endsWith(".jar"));
            if (jarFiles == null) {
                throw new EdsException("目录下没有jar文件!");
            }
            return Arrays.stream(jarFiles).map(file -> {
                try {
                    return new URL(toJarUrl(file.getAbsolutePath()));
                } catch (MalformedURLException e) {
                    throw new EdsException(e);
                }
            }).toArray(URL[]::new);
        }
        try {
            return new URL[]{new URL(toJarUrl(jarPath))};
        } catch (MalformedURLException e) {
            throw new EdsException(e);
        }
    }

    private String toJarUrl(String url) {
        return "jar:file:" + url + "!/";
    }

    private void closeAllDataSources() {
        // 关闭静态数据源
        for (ExpireDataSourceHolder holder : staticDataSources.values()) {
            try {
                holder.close();
            } catch (Exception e) {
                log.warn("关闭静态数据源失败: {}", holder.getDsKey(), e);
            }
        }

        // 关闭目标数据源
        for (ExpireDataSourceHolder holder : targetDataSources.values()) {
            try {
                holder.close();
            } catch (Exception e) {
                log.warn("关闭目标数据源失败: {}", holder.getDsKey(), e);
            }
        }

        // 清空缓存
        staticDataSources.clear();
        targetDataSources.clear();
        evictQueue.clear();
        classLoaderMap.clear();
    }

    public final class DestroyEdsThread extends Thread {
        public DestroyEdsThread(String name) {
            super(name);
        }

        @Override
        public void run() {
            // 启动后延迟执行
            long sleepTime = Math.max(staticDataSourceProperties.getEdsScanTime(), 100);
            log.info("销毁数据源线程启动,调度间隔{}ms", sleepTime);
            //noinspection InfiniteLoopStatement
            for (; ; ) {
                try {
                    //noinspection BusyWait
                    Thread.sleep(sleepTime);
                    if (evictQueue.isEmpty()) {
                        continue;
                    }

                    int closedCount = 0;
                    // 处理队列中的数据源
                    List<ExpireDataSourceHolder> remaining = null;
                    // 先取出所有待处理的数据源
                    ExpireDataSourceHolder holder;
                    while ((holder = evictQueue.poll()) != null) {
                        if (holder.isCanClose()) {
                            try {
                                holder.close();
                                closedCount++;
                            } catch (Exception e) {
                                log.warn("关闭数据源异常，继续关闭下一个!dsKey={},ex={}", holder.getDsKey(), e);
                            }
                        } else {
                            if (remaining == null) {
                                remaining = new ArrayList<>();
                            }
                            remaining.add(holder);
                        }
                    }

                    // 将无法关闭的数据源重新放回队列
                    if (remaining != null) {
                        for (ExpireDataSourceHolder remainingHolder : remaining) {
                            evictQueue.put(remainingHolder);
                        }
                    }
                    if (closedCount > 0) {
                        log.info("清理数据源={},剩余数据源={},evictQueue总量={}", closedCount, targetDataSources.size(), evictQueue.size());
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    log.warn("eds销毁线程中断，准备退出!");
                    return;
                } catch (Exception e) {
                    log.info("eds销毁线程出现异常!", e);
                }
            }
        }
    }

    public class ScanExpiredEdsThread extends Thread {
        public ScanExpiredEdsThread(String name) {
            super(name);
        }

        @Override
        public void run() {
            // 启动后延迟执行
            long sleepTime = Math.max(staticDataSourceProperties.getEdsScanTime(), 100);
            log.info("扫描过期数据源线程启动,扫描间隔{}ms", sleepTime);
            //noinspection InfiniteLoopStatement
            for (; ; ) {
                try {
                    //noinspection BusyWait
                    Thread.sleep(sleepTime);
                    if (targetDataSources.isEmpty()) {
                        continue;
                    }
                    long now = System.currentTimeMillis();
                    List<ExpireDataSourceHolder> evictList = new ArrayList<>(targetDataSources.size());
                    // 标记出所有要销毁的连接池
                    for (Iterator<Map.Entry<String, ExpireDataSourceHolder>> it = targetDataSources.entrySet().iterator(); it.hasNext(); ) {
                        Map.Entry<String, ExpireDataSourceHolder> entry = it.next();
                        ExpireDataSourceHolder dsHolder = entry.getValue();
                        // accessTime不要求获取最新的值（稍微滞后可以接受）
                        if (dsHolder != null && isExpired(dsHolder.getAccessTime(), now)) {
                            evictList.add(dsHolder);
                            dsHolder.setRemoved(true);
                            it.remove();
                        }
                    }
                    // 批量添加到队列
                    for (ExpireDataSourceHolder holder : evictList) {
                        evictQueue.put(holder);
                    }
                    long endTime = System.currentTimeMillis();
                    log.info("扫描过期数据源总量={},扫描耗时={}ms,evictQueue总量={}", evictList.size(), endTime - now, evictQueue.size());
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    log.warn("eds扫描线程中断，准备退出!");
                    return;
                } catch (Exception e) {
                    log.warn("scan datasource thread error!", e);
                }
            }
        }

        private boolean isExpired(long accessTime, long now) {
            return now - accessTime > staticDataSourceProperties.getEdsExpireTime();
        }
    }
}
